//
// Created by Peilin Fan
//

#include <iostream>
#include <string>

#include "gpt2.h"

/**
 * @brief GPT-2模型的结束符token ID
 */
constexpr int GPT2_EOT = 50256;

/**
 * @brief 主函数
 * @param argc 命令行参数数量
 * @param argv 命令行参数数组
 * @return 程序退出状态码
 * 
 * 程序功能：
 * 1. 从命令行接收初始token序列
 * 2. 使用GPT-2模型生成后续token
 * 3. 每生成一个token就输出到标准输出
 * 
 * 使用方式：
 * ./gpt2 token1 token2 token3 ...
 */
int main(int argc, char** argv) {
    // 检查命令行参数
    if (argc < 3) {
        std::cerr << "Provide at least one token." << std::endl;
        std::terminate();
    }
    // 设置生成序列的最大长度
    const int N = std::stoi(argv[1]);
    if (argc > N + 1) {
        std::cerr << "Too many tokens." << std::endl;
        std::terminate();
    }

    // 初始化GPT-2模型
    GPT2 model{"gpt2_124M.bin"};
    // 初始化token序列数组
    int tokens[up8(N)];

    // 填充token序列：
    // 1. 先填入命令行提供的token
    // 2. 剩余位置填充结束符
    for (int i = 0; i < up8(N); i++) {
        if (i + 2 < argc) {
            tokens[i] = std::stoi(argv[i + 2]);
        } else {
            tokens[i] = GPT2_EOT;
        }
    }

    // 从已有序列开始，逐个生成新token
    for (int t = argc - 2; t < N; t++) {
        // 使用模型生成下一个token
        tokens[t] = model.generate_one(tokens, t);

        // 输出生成的token并刷新缓冲区
        std::cout << tokens[t] << std::endl;
        std::cout.flush();
    }
}
